// SPDX-FileCopyrightText: Copyright 2025-2025 深圳市同心圆网络有限公司
// SPDX-License-Identifier: GPL-3.0-only

package badger_dao_impl

import (
	"context"
	"encoding/json"
	"path/filepath"
	"sort"
	"strings"

	"gitcode.com/openseaotter/so_server/storage/core"
	"gitcode.com/openseaotter/so_server/utils"
	"github.com/dgraph-io/badger/v4"
)

type _ImageReferenceDao struct {
	dbh *badger.DB
}

func NewImageReferenceDao() (core.ImageReferenceDao, error) {
	dbPath := filepath.Join(utils.DATA_DIR_PATH, "image_reference")
	options := badger.DefaultOptions(dbPath)
	dbh, err := badger.Open(options)
	if err != nil {
		return nil, err
	}
	return &_ImageReferenceDao{dbh: dbh}, nil
}

func (db *_ImageReferenceDao) genKey(groupName, imageName, reference string) string {
	return strings.Join([]string{groupName, imageName, reference}, "\t")
}

func (db *_ImageReferenceDao) genKeyPrefix(groupName, imageName string) string {
	return strings.Join([]string{groupName, imageName}, "\t") + "\t"
}

func (db *_ImageReferenceDao) Insert(ctx context.Context, imageRefInfo *core.ImageReferenceInfo) error {
	return db.Update(ctx, imageRefInfo)
}

func (db *_ImageReferenceDao) Exist(ctx context.Context, groupName, imageName, reference string) (bool, error) {
	txn := db.dbh.NewTransaction(false)
	defer txn.Discard()

	key := db.genKey(groupName, imageName, reference)
	item, err := txn.Get([]byte(key))
	if err != nil {
		if err == badger.ErrKeyNotFound {
			return false, nil
		}
		return false, err
	}
	if item.IsDeletedOrExpired() {
		return false, nil
	}
	return true, nil
}

func (db *_ImageReferenceDao) Get(ctx context.Context, groupName, imageName, reference string) (*core.ImageReferenceInfo, error) {
	txn := db.dbh.NewTransaction(false)
	defer txn.Discard()

	key := db.genKey(groupName, imageName, reference)
	item, err := txn.Get([]byte(key))
	if err != nil {
		return nil, err
	}
	if item.IsDeletedOrExpired() {
		return nil, badger.ErrKeyNotFound
	}
	value, err := item.ValueCopy(nil)
	if err != nil {
		return nil, err
	}
	info := &core.ImageReferenceInfo{}
	err = json.Unmarshal(value, info)
	if err != nil {
		return nil, err
	}
	return info, nil
}

func (db *_ImageReferenceDao) Update(ctx context.Context, imageRefInfo *core.ImageReferenceInfo) error {
	txn := db.dbh.NewTransaction(true)
	defer txn.Discard()

	key := db.genKey(imageRefInfo.GroupName, imageRefInfo.ImageName, imageRefInfo.Reference)
	value, err := json.Marshal(imageRefInfo)
	if err != nil {
		return err
	}

	err = txn.Set([]byte(key), value)
	if err != nil {
		return err
	}
	return txn.Commit()
}

func (db *_ImageReferenceDao) Remove(ctx context.Context, groupName, imageName, reference string) error {
	txn := db.dbh.NewTransaction(true)
	defer txn.Discard()

	key := db.genKey(groupName, imageName, reference)

	err := txn.Delete([]byte(key))
	if err != nil {
		return err
	}
	return txn.Commit()
}

func (db *_ImageReferenceDao) RemoveByImage(ctx context.Context, groupName, imageName string) error {
	txn := db.dbh.NewTransaction(true)
	defer txn.Discard()

	keyPrefix := db.genKeyPrefix(groupName, imageName)
	options := badger.IteratorOptions{
		PrefetchValues: false,
		PrefetchSize:   100,
		Reverse:        false,
		AllVersions:    false,
		Prefix:         []byte(keyPrefix),
	}
	iter := txn.NewIterator(options)
	defer iter.Close()

	keyList := [][]byte{}
	for iter.Rewind(); iter.Valid(); iter.Next() {
		item := iter.Item()
		if item.IsDeletedOrExpired() {
			continue
		}
		key := item.KeyCopy(nil)
		keyList = append(keyList, key)
	}

	iter.Close()

	for _, key := range keyList {
		err := txn.Delete(key)
		if err != nil {
			return err
		}
	}
	return txn.Commit()
}

func (db *_ImageReferenceDao) Count(ctx context.Context, groupName, imageName string, includeDigest bool) (uint32, error) {
	txn := db.dbh.NewTransaction(false)
	defer txn.Discard()

	keyPrefix := db.genKeyPrefix(groupName, imageName)
	options := badger.IteratorOptions{
		PrefetchValues: false,
		PrefetchSize:   100,
		Reverse:        false,
		AllVersions:    false,
		Prefix:         []byte(keyPrefix),
	}
	iter := txn.NewIterator(options)
	defer iter.Close()

	count := uint32(0)

	for iter.Rewind(); iter.Valid(); iter.Next() {
		item := iter.Item()
		if item.IsDeletedOrExpired() {
			continue
		}
		key := item.KeyCopy(nil)
		parts := strings.Split(string(key), "\t")
		if len(parts) != 3 {
			continue
		}
		if strings.HasPrefix(parts[2], "sha256:") || strings.HasPrefix(parts[2], "sha384:") || strings.HasPrefix(parts[2], "sha512:") {
			if !includeDigest {
				continue
			}
		}
		count += 1
	}

	return count, nil
}

func (db *_ImageReferenceDao) List(ctx context.Context, groupName, imageName string, includeDigest bool) ([]*core.ImageReferenceInfo, error) {
	txn := db.dbh.NewTransaction(false)
	defer txn.Discard()

	keyPrefix := db.genKeyPrefix(groupName, imageName)
	options := badger.IteratorOptions{
		PrefetchValues: true,
		PrefetchSize:   100,
		Reverse:        false,
		AllVersions:    false,
		Prefix:         []byte(keyPrefix),
	}
	iter := txn.NewIterator(options)
	defer iter.Close()

	infoList := []*core.ImageReferenceInfo{}
	for iter.Rewind(); iter.Valid(); iter.Next() {
		item := iter.Item()
		if item.IsDeletedOrExpired() {
			continue
		}
		key := item.KeyCopy(nil)
		parts := strings.Split(string(key), "\t")
		if len(parts) != 3 {
			continue
		}
		if strings.HasPrefix(parts[2], "sha256:") || strings.HasPrefix(parts[2], "sha384:") || strings.HasPrefix(parts[2], "sha512:") {
			if !includeDigest {
				continue
			}
		}
		value, err := item.ValueCopy(nil)
		if err != nil {
			return nil, err
		}
		info := &core.ImageReferenceInfo{}
		err = json.Unmarshal(value, info)
		if err != nil {
			return nil, err
		}
		infoList = append(infoList, info)
	}

	sort.Sort(sort.Reverse(core.ImageReferenceInfoList(infoList)))
	return infoList, nil
}

func (db *_ImageReferenceDao) ListAll(ctx context.Context) ([]*core.ImageReferenceInfo, error) {
	txn := db.dbh.NewTransaction(false)
	defer txn.Discard()

	options := badger.IteratorOptions{
		PrefetchValues: true,
		PrefetchSize:   100,
		Reverse:        false,
		AllVersions:    false,
	}
	iter := txn.NewIterator(options)
	defer iter.Close()

	infoList := []*core.ImageReferenceInfo{}
	for iter.Rewind(); iter.Valid(); iter.Next() {
		item := iter.Item()
		if item.IsDeletedOrExpired() {
			continue
		}
		value, err := item.ValueCopy(nil)
		if err != nil {
			return nil, err
		}
		info := &core.ImageReferenceInfo{}
		err = json.Unmarshal(value, info)
		if err != nil {
			return nil, err
		}
		infoList = append(infoList, info)
	}
	return infoList, nil
}
